# coding=utf8
from __future__ import division
import torch
import os, time, datetime
from torch.autograd import Variable
import logging
import torch.nn.functional as F
import numpy as np
from math import ceil
import copy
from predicting import get_class_recs, voc_ap, cal_class_ap
from plotting import colors, plot_img
import matplotlib.pyplot as plt
import cv2

label_name_map = {0:'defect'}

def dt():
    return datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')


def trainlog(logfilepath, head='%(message)s'):
    logger = logging.getLogger('mylogger')
    logging.basicConfig(filename=logfilepath, level=logging.INFO, format=head)

    console = logging.StreamHandler()
    console.setLevel(logging.INFO)
    formatter = logging.Formatter(head)
    console.setFormatter(formatter)
    logging.getLogger('').addHandler(console)


def train(model,
          epoch_num,
          start_epoch,
          optimizer,
          criterion,
          exp_lr_scheduler,
          data_set,
          data_loader,
          save_dir,
          print_inter=200,
          val_inter=2000,
          ):
    best_model_wts = model.state_dict()
    best_mAP = 0.0

    step = -1
    running_loss = 30
    model_num_classes = model.module.num_classes

    for epoch in range(start_epoch, epoch_num):
        # train phase
        exp_lr_scheduler.step(epoch)
        model.train(True)  # Set model to training mode

        for batch_cnt, data in enumerate(data_loader['train']):

            step += 1
            model.train(True)
            model.module.phase = 'train'
            # print data
            inputs, targets, patch_pos = data

            inputs = Variable(inputs.cuda())
            targets = [Variable(anno.cuda()) for anno in targets]



            # zero the parameter gradients
            optimizer.zero_grad()

            out = model(inputs)
            loss_l, loss_c = criterion(out, targets)
            loss = loss_l + loss_c

            running_loss = running_loss*0.9 + 0.1*loss.data[0]

            loss.backward()
            optimizer.step()

            # batch loss
            if step % print_inter == 0:
                logging.info('%s [%d]-[%d] || batch-loss: %.4f || running-loss: %.4f'
                             % (dt(), epoch, step, loss.data[0], running_loss))

            #


            # eval model
            if step % val_inter == 0:
                model.module.phase = 'test'

                idx_val = 0
                val_patch_pos = np.empty((len(data_set['val']), 6), dtype=np.int)
                det_result = np.zeros((len(data_set['val']), model_num_classes, 50, 5), dtype=np.float32)

                for batch_cnt, data in enumerate(data_loader['val']):
                    # if batch_cnt == 5:break
                    inputs, targets, patch_pos = data  # targets [arr(n1,5), arr(n2,5), ...]
                    inputs = Variable(inputs.cuda())

                    detections = model(inputs)
                    detections = detections.cpu().data.numpy()  # (N,1+C,50,5) confidence first

                    det_result[idx_val: idx_val+detections.shape[0]] = detections   # float localization
                    val_patch_pos[idx_val: idx_val+detections.shape[0]] = patch_pos
                    idx_val = idx_val+detections.shape[0]

                    ##############
                    if step>0:
                        for i in range(len(data_set['val'])):
                            img_path = data_set['val'].img_paths[i]
                            input_img = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
                            input_img = cv2.resize(input_img, (300, 300))
                            annos = data_set['val'].bboxs_pos[i]

                            input_img, img_true, img_pred, _, _ = plot_img(model, input_img, annos=annos, transforms=None)
                            plt.figure()
                            plt.subplot(131)
                            plt.imshow(input_img)
                            plt.subplot(132)
                            plt.imshow(img_true)
                            plt.subplot(133)
                            plt.imshow(img_pred)
                            plt.show()

                # float localization to absolute
                det_result[:, :, :, 1] = det_result[:, :, :, 1] * val_patch_pos[:, 4][:,np.newaxis,np.newaxis]
                det_result[:, :, :, 3] = det_result[:, :, :, 3] * val_patch_pos[:, 4][:,np.newaxis,np.newaxis]
                det_result[:, :, :, 2] = det_result[:, :, :, 2] * val_patch_pos[:, 5][:,np.newaxis,np.newaxis]
                det_result[:, :, :, 4] = det_result[:, :, :, 4] * val_patch_pos[:, 5][:,np.newaxis,np.newaxis]


                det_ids = np.arange(det_result.shape[0]).reshape((-1,1,1,1))
                det_ids = np.ones(list(det_result.shape[0:3])+[1])*det_ids

                det_result = np.concatenate((det_ids, det_result), axis=3)  # (num_imgs, classes, 200, 6)   6:[id,confi, xmin,ymin,xmax,ymax]

                mAP = []
                for class_name in data_set['val'].class_map.keys():

                    # extract gt objects for this class
                    class_recs, npos, class_idx = get_class_recs(data_set['val'], class_name=class_name)

                    # extract det for this class
                    class_dets = det_result[:,class_idx,:,:].reshape(-1,6)   # [id, confi, xmin, ymin, xmax, ymax]
                    class_dets = class_dets[class_dets[:,1]>0.3]  # detections with confidence greater than 0.0


                    # cal AP for this class
                    rec, prec, ap = cal_class_ap(class_dets, class_recs, npos, ovthresh = 0.3)
                    mAP.append(ap)
                    logging.info('AP for "%s" : %.4f' % (class_name, ap))

                mAP = 1.*sum(mAP)/len(mAP)
                logging.info('step-%d mAP: %.4f' % (step, mAP))

                # save model
                if mAP > best_mAP:
                    best_mAP = mAP
                    best_model_wts = copy.deepcopy(model.state_dict())

                save_path = os.path.join(save_dir,
                                         'weights-%d-%d-[%.4f].pth' % (epoch, step, mAP))
                torch.save(model.state_dict(), save_path)
                logging.info('saved model to %s' % (save_path))
                logging.info('--' * 30)
    #
    # save best model
    save_path = os.path.join(save_dir,
                             'bestweights-[%.4f].pth' % ( best_mAP))
    torch.save(best_model_wts, save_path)
    logging.info('saved model to %s' % (save_path))

    return best_mAP, best_model_wts
